# -*- coding:utf-8

import numpy as np
from scipy.special import gamma
from .formula import *

__all__ = ['WindStaterBase', 'DefaultWindStater', 'wind_stater']


class WindStaterBase(object):
    WSBINS = np.arange(1.0, 50.1, 1.0, dtype='f4')   # 风速段
    WSBINSTRS = np.array(['%g' % v for v in WSBINS])  # 字符串表示的风速段
    N_WD = 16   # 风向扇区数
    WD0 = 0.0   # 第一个风向扇区的中心角度
    CALM_THRES = 0.3  # 静风风速阈值 (m/s)
    N_WS = len(WSBINS)   # 风速段数
    N_WDC = N_WD + 1    # 风向扇区数(带上静风)
    WDBINS = np.fmod((np.arange(N_WD) * (360.0 / N_WD) + WD0), 360.0)  # (风向扇区)

    # 将风速数组的值转换为在风速段中的下标
    # 该函数需要随WSBINS的不同而手动修改, 以实现高效的转换
    def ws2wsbin(self, ws):
        res = ws.astype('i4')
        res[res>49] = 49
        res[res<0] = -1
        return res

    def wd2wdbin(self, wd):
        binsize = 360.0 / self.N_WD
        offsetted_wd = wd + binsize/2.0 - self.WD0 + 360.0  # TODO 检查 WD0 是否合理
        return np.array(offsetted_wd / binsize).astype('i4') % self.N_WD

    def return_as_needed(self, count, total_count, return_type='percent'):
        if return_type == 'count':
            return count
        else:
            ratio = count / float(total_count)
            if return_type == 'ratio':
                return ratio
            elif return_type == 'permill':
                return ratio * 1000.0
            else:
                return ratio * 100.0


    def get_wind_rose(self, ws, wd, return_type='percent'):
        """
        :param ws: 1维的风速数组
        :param wd: 1维的风向数组
        :param return_type: 返回类型: 'count': 数量 'ratio': 小数值, 'percent': 百分比, 'permill': 千分比
        :return:
        """
        stacker = np.zeros((len(ws), self.N_WS, self.N_WD), dtype='i4')
        ws_binned = self.ws2wsbin(ws)
        wd_binned = self.wd2wdbin(wd)
        stacker[(np.arange(len(ws)), ws_binned, wd_binned)] = 1
        stacker[np.where(ws_binned<0)] = 0  # 对付掉nan导致的ws_binned中的负值
        squeezed = np.sum(stacker, axis=0)

        return self.return_as_needed(squeezed, np.ma.masked_invalid(ws).count(), return_type=return_type)


    def get_wind_rose_2d(self, ws, wd, return_type='percent'):
        res_dtype = 'i4' if return_type == 'count' else 'f4'
        N_JY = ws.shape[1]
        N_IX = ws.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WS, self.N_WD), dtype=res_dtype)
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wind_rose(ws[:, jy, ix], wd[:, jy, ix], return_type=return_type)
        return res_array


    def get_wpd_rose(self, ws, wd, wpd=None, rhoair=1.125):
        stacker = np.zeros((len(ws), self.N_WS, self.N_WD), dtype='f8')
        if wpd is None:
            wpd = calc_wpd(ws, rhoair=rhoair)
        ws_binned = self.ws2wsbin(ws)
        wd_binned = self.wd2wdbin(wd)
        stacker[(np.arange(len(ws)), ws_binned, wd_binned)] = 1.0
        stacker[np.where(ws_binned<0)] = 0
        stacker *= wpd[:, np.newaxis, np.newaxis]
        res = np.nansum(stacker, axis=0)
        return res


    def get_wpd_rose_2d(self, ws, wd, wpd=None, rhoair=1.125):
        N_JY = ws.shape[1]
        N_IX = ws.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WS, self.N_WD), dtype='f4')
        for jy, ix in np.ndindex(N_JY, N_IX):
            if wpd is None:
                wpd_one = None
            else:
                wpd_one = wpd[:, jy, ix]
            # todo: handle rhoair array
            res_array[jy, ix] = self.get_wpd_rose(ws[:, jy, ix], wd[:, jy, ix], wpd=wpd_one, rhoair=rhoair)
        return res_array


    def get_wsdist(self, ws, return_type='percent'):
        stacker = np.zeros((len(ws), self.N_WS), dtype='i4')
        ws_binned = self.ws2wsbin(ws)
        stacker[(np.arange(len(ws)), ws_binned)] = 1
        stacker[np.where(ws_binned<0)] = 0
        squeezed = np.sum(stacker, axis=0)
        return self.return_as_needed(squeezed, np.ma.masked_invalid(ws).count(), return_type=return_type)

    def get_wsdist_2d(self, ws, return_type='percent'):
        N_JY = ws.shape[1]
        N_IX = ws.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WS), dtype='f4')
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wsdist(ws[:, jy, ix], return_type=return_type)
        return res_array

    def get_wddist(self, wd, return_type='percent'):
        stacker = np.zeros((len(wd), self.N_WD), dtype='i4')
        wd_binned = self.wd2wdbin(wd)
        stacker[(np.arange(len(wd)), wd_binned)] = 1
        stacker[np.where(~np.isfinite(wd))] = 0
        squeezed = np.sum(stacker, axis=0)
        return self.return_as_needed(squeezed, np.ma.masked_invalid(wd).count(), return_type=return_type)

    def get_wddist_2d(self, wd, return_type='percent'):
        N_JY = wd.shape[1]
        N_IX = wd.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WD), dtype='f4')
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wddist(wd[:, jy, ix], return_type=return_type)
        return res_array

    def get_wdmean(self, wd, ws):
        """各扇区的平均风速"""
        stacker = np.zeros((len(wd), self.N_WD), dtype='f4')
        stacker[:] = np.nan
        wd_binned = self.wd2wdbin(wd)
        stacker[(np.arange(len(wd)), wd_binned)] = 1.0
        stacker[np.where(~np.isfinite(wd))] = 0.0
        stacker *= ws[:, np.newaxis]
        squeezed = np.nanmean(stacker, axis=0)
        return squeezed

    def get_wdmean_2d(self, wd, ws):
        N_JY = wd.shape[1]
        N_IX = wd.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WD), dtype='f4')
        res_array[:] = np.nan
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wdmean(wd[:, jy, ix], ws[:, jy, ix])
        return res_array

    def get_wdcdist(self, wd, ws, return_type='percent'):
        w_calm = np.where(ws < self.CALM_THRES)
        w_active = np.where(~(ws < self.CALM_THRES))  # 把nan归到这一类里, 让get_wddist处理
        wd_active = wd[w_active]
        wddist_count = self.get_wddist(wd_active, return_type='count')
        calm_count = len(w_calm[0])
        wdcdist_count = np.zeros(self.N_WDC, 'i4')
        wdcdist_count[:self.N_WD] = wddist_count
        wdcdist_count[-1] = calm_count
        return self.return_as_needed(wdcdist_count, np.ma.masked_invalid(wd).count(), return_type=return_type)

    def get_wdcdist_2d(self, wd, ws, return_type='percent'):
        N_JY = wd.shape[1]
        N_IX = wd.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WDC), dtype='f4')
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wdcdist(wd[:, jy, ix], ws[:, jy, ix], return_type=return_type)
        return res_array

    def get_ak_gamma(self, wsdist, dist_type='percent'):
        """
        由王瑞明提供的方法
        :param wsdist:
        :param dist_type:
        :return:
        """
        if dist_type == 'raw':
            ws = wsdist
            ws = np.ma.masked_invalid(ws)
            k = (ws.std() / ws.mean()) ** -1.086
            a = ((ws*ws*ws).mean() / gamma(3.0 / k + 1.0)) ** (1.0/3)
        else:
            if dist_type == 'count':
                wsdist_ratio = wsdist / float(wsdist.sum())
            elif dist_type == 'percent':
                wsdist_ratio = wsdist / 100.0
            elif dist_type == 'permill':
                wsdist_ratio = wsdist / 1000.0
            else:
                wsdist_ratio = wsdist
            WSBINS_0 = np.concatenate(([0.0], self.WSBINS))
            wsmid = (WSBINS_0[:-1] + WSBINS_0[1:]) / 2.0
            ws_mean = np.sum(wsdist_ratio * wsmid)
            ws_std = np.sqrt(np.sum(wsdist_ratio * np.square(wsmid - ws_mean)))
            ws3_mean = np.sum(wsdist_ratio * wsmid * wsmid * wsmid)
            k = (ws_std / ws_mean) ** -1.086
            a = (ws3_mean / gamma(3.0 / k + 1.0)) ** (1.0/3)
        return a, k

    def get_wpd_wsdist(self, wpd, ws):
        stacker = np.zeros((len(ws), self.N_WS), dtype='f8')
        ws_binned = self.ws2wsbin(ws)
        stacker[(np.arange(len(ws)), ws_binned)] = 1.0
        stacker[np.where(ws_binned<0)] = 0.0
        stacker *= wpd[:, np.newaxis]
        squeezed = np.nansum(stacker, axis=0)
        return squeezed

    def get_wpd_wsdist_2d(self, wpd, ws):
        N_JY = ws.shape[1]
        N_IX = ws.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WS), dtype='f4')
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wpd_wsdist(wpd[:, jy, ix], ws[:, jy, ix])
        return res_array

    def get_wpd_wddist(self, wpd, wd):
        stacker = np.zeros((len(wd), self.N_WD), dtype='f8')
        wd_binned = self.wd2wdbin(wd)
        stacker[(np.arange(len(wd)), wd_binned)] = 1.0
        stacker[np.where(~np.isfinite(wd))] = 0.0
        stacker *= wpd[:, np.newaxis]
        squeezed = np.nansum(stacker, axis=0)
        return squeezed

    def get_wpd_wddist_2d(self, wpd, wd):
        N_JY = wd.shape[1]
        N_IX = wd.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WD), dtype='f4')
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wpd_wddist(wpd[:, jy, ix], wd[:, jy, ix])
        return res_array

    def get_wpd_wdcdist(self, wpd, wd, ws):
        w_calm = np.where(ws < self.CALM_THRES)
        w_active = np.where(~(ws < self.CALM_THRES))  # 把nan归到这一类里, 让get_wddist处理
        wd_active = wd[w_active]
        wpd_active = wpd[w_active]
        wpd_calm = wpd[w_calm]
        wpd_wddist_active = self.get_wpd_wddist(wpd_active, wd_active)
        wpd_wddist_calm = np.nansum(wpd_calm)
        wdcdist = np.zeros(self.N_WDC, 'f4')
        wdcdist[:self.N_WD] = wpd_wddist_active
        wdcdist[-1] = wpd_wddist_calm
        return wdcdist

    def get_wpd_wdcdist_2d(self, wpd, wd, ws):
        N_JY = wd.shape[1]
        N_IX = wd.shape[2]
        res_array = np.zeros((N_JY, N_IX, self.N_WDC), dtype='f4')
        for jy, ix in np.ndindex(N_JY, N_IX):
            res_array[jy, ix] = self.get_wpd_wdcdist(wpd[:, jy, ix], wd[:, jy, ix], ws[:, jy, ix])
        return res_array


class DefaultWindStater(WindStaterBase):
    pass


wind_stater = DefaultWindStater()